//
//  RBT.h
//  RBT
//
//  Created by wushiyu on 2021/1/24.
//

#ifndef RBT_h
#define RBT_h
using namespace std;
template <typename T>
class Node
{
public:
    T data;
    char color;
    Node<T>* lc;
    Node<T>* rc;
    Node<T>* parent;
    Node(T _data);
};

template <typename T>
class RBT{
public:
    Node<T>* root;
    RBT(){root = NULL;}
    RBT(T _data);
    ~RBT(){};
    Node<T>* minimum(Node<T>* x);
    void left_rotate(Node<T>* x);
    void right_rotate(Node<T>* x);
    void Insert(T _data);
    void Insert_(Node<T>* z);
    void Delete(Node<T>* z);
    void transplant(Node<T>* u, Node<T>* v);
    void Delete_(Node<T>* x);
};

template <typename T>
Node<T>::Node(T _data){
    data = _data;
    parent = lc = rc = NULL;
}

template <typename T>
RBT<T>::RBT(T _data){
    Node<T>* m = new Node<T>(_data);
    if (root == NULL)
    root = m;
    root->color = 'b';
};

template <typename T>
void RBT<T>::left_rotate(Node<T>* x){
    Node<T>* y = x->rc;
    x->rc = y->lc;
    if(y->lc != NULL)
        y->lc->parent = x;
    y->parent = x->parent;
    if(x->parent == NULL)
        root = y;
    else if(x == x->parent->lc)
        x->parent->lc = y;
    else x->parent->rc = y;
    y->lc = x;
    x->parent = y;
}

template <typename T>
void RBT<T>::right_rotate(Node<T>* x){
    Node<T>* y = x->lc;
    x->lc = y->rc;
    if(y->rc !=  NULL)
        y->rc->parent = x;
    y->parent = x->parent;
    if(x->parent == NULL)
        root = y;
    else if(x == x->parent->lc)
        x->parent->rc = y;
    else x->parent->lc = y;
    y->rc = x;
    x->parent = y;
}

template <typename T>
void RBT<T>::Insert(T _data){
    Node<T>* z = new Node<T>(_data);
    Node<T>* y = NULL;
    Node<T>* x = root;
    while(x != NULL){
        y = x;
        if(z->data < x->data)
            x = x->lc;
        else x = x->rc;
    }
    z->parent = y;
    if(y == NULL)
        root = z;
    else if(z->data < y->data)
        y->lc = z;
    else y->rc = z;
    z->lc = NULL;
    z->rc = NULL;
    z->color = 'r';
    RBT<T>::Insert_(z);
    root->color = 'b';
}

template <typename T>
void RBT<T>::Insert_(Node<T>* z){
    while(z->parent != NULL && z->parent->color == 'r'){
        if(z->parent == z->parent->parent->lc){
            Node<T>* y = z->parent->parent->rc;
            if(y != NULL && y->color == 'r'){
                z->parent->color = 'b';
                y->color = 'b';
                z->parent->parent->color = 'r';
                z = z->parent->parent;
            }
            else if(z != NULL && z == z->parent->rc){
                z = z->parent;
                RBT<T>::left_rotate(z);
            }
            if(z->parent->parent != NULL){
                z->parent->color = 'b';
                z->parent->parent->color = 'r';
                RBT<T>::right_rotate(z->parent->parent);
            }
        }
        else{
            Node<T>* y = z->parent->parent->lc;
            if(y != NULL && y->color == 'r'){
                z->parent->color = 'b';
                y->color = 'b';
                z->parent->parent->color = 'r';
                z = z->parent->parent;
            }
            else if(z == z->parent->lc){
                z = z->parent;
                RBT<T>::right_rotate(z);
            }
            z->parent->color = 'b';
            z->parent->parent->color = 'r';
            RBT<T>::left_rotate(z->parent->parent);
        }
    }
}
template <typename T>
Node<T>* RBT<T>::minimum(Node<T>* x)
{
    while (x->lc != NULL)
    x = x->lc;
    return x;
};

template <typename T>
void RBT<T>::transplant(Node<T> *u, Node<T> *v){
    if(u->parent == NULL)
        root = v;
    else if(u == u->parent->lc)
        u->parent->lc = v;
    else u->parent->rc = v;
    if(v != NULL)
        v->parent = u->parent;
}
template <typename T>
void RBT<T>::Delete(Node<T> *z){
    Node<T>* y = z;
    Node<T>* x;
    char yoc = y->color;
    if(z->lc == NULL){
        x = z->rc;
        RBT<T>::transplant(z, z->rc);
    }
    else if(z->rc == NULL){
        x = z->lc;
        RBT<T>::transplant(z, z->lc);
    }
    else{
        y = RBT<T>::minimum(z->rc);
        yoc = y->color;
        x = y->rc;
    }
    if(x != NULL){
        if(y->parent == z)
            x->parent = y;
        else{
            RBT<T>::transplant(y, y->rc);
            y->rc = z->rc;
            y->rc->parent = y;
        }
    }
    RBT<T>::transplant(z, y);
    y->lc = z->lc;
    y->lc->parent = y;
    y->color = z->color;
    if(yoc == 'b')
        RBT<T>::Delete_(x);
}
template <typename T>
void RBT<T>::Delete_(Node<T> *x){
    while(x != root && x->color == 'b'){
        if(x == x->parent->lc){
            Node<T>* w = x->parent->rc;
            if(w->color == 'r'){
                w->color = 'b';
                x->parent->color = 'r';
                RBT<T>::left_rotate(x->parent);
                w = x->parent->rc;
            }
            if(w->lc->color == 'b' && w->rc->color == 'b'){
                w->color = 'r';
                x = x->parent;
            }
            else if(w->rc->color == 'b'){
                w->lc->color = 'b';
                w->color = 'r';
                RBT<T>::right_rotate(w);
                w = x->parent->rc;
            }
            w->color = x->parent->color;
            x->parent->color = 'b';
            w->rc->color = 'b';
            RBT<T>::left_rotate(x->parent);
            x = root;
        }
        else{
            if(x == x->parent->rc){
                Node<T>* w = x->parent->lc;
                if(w->color == 'r'){
                    w->color = 'b';
                    x->parent->color = 'r';
                    RBT<T>::right_rotate(x->parent);
                    w = x->parent->lc;
                }
                if(w->rc->color == 'b' && w->lc->color == 'b'){
                    w->color = 'r';
                    x = x->parent;
                }
                else if(w->lc->color == 'b'){
                    w->rc->color = 'b';
                    w->color = 'r';
                    RBT<T>::left_rotate(w);
                    w = x->parent->lc;
                }
                w->color = x->parent->color;
                x->parent->color = 'b';
                w->lc->color = 'b';
                RBT<T>::right_rotate(x->parent);
                x = root;
            }
        }
    }
    x->color = 'b';
}
#endif /* RBT_h */
